import os
import time
import uuid
import json
import base64
import sqlite3
from datetime import datetime, timedelta
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding

class KeyManager:
    """
    密钥生命周期管理类
    实现密钥的生成、存储、轮换、销毁等功能
    """
    
    def __init__(self, db_path=None):
        # 使用SQLite而不是MySQL
        self.db_path = db_path or os.path.join(os.path.dirname(os.path.dirname(__file__)), 'crypto_db.sqlite')
        self._connection = None
        self._ensure_db_schema()
    
    def _get_connection(self):
        """获取数据库连接"""
        if self._connection is None:
            self._connection = sqlite3.connect(self.db_path)
            # 启用外键约束
            self._connection.execute("PRAGMA foreign_keys = ON")
            # 将SQLite设置为返回行类似于字典
            self._connection.row_factory = sqlite3.Row
        return self._connection
    
    def _ensure_db_schema(self):
        """确保数据库表结构存在"""
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            # 创建密钥表
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS `keys` (
                `id` TEXT NOT NULL PRIMARY KEY,
                `key_name` TEXT NOT NULL,
                `key_type` TEXT NOT NULL,
                `key_data` TEXT NOT NULL,
                `salt` TEXT NULL,
                `created_at` TIMESTAMP NOT NULL,
                `expires_at` TIMESTAMP NULL,
                `revoked_at` TIMESTAMP NULL,
                `metadata` TEXT NULL
            );
            """)
            
            # 创建索引
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_key_name ON keys (key_name);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_created_at ON keys (created_at);")
            
            # 创建密钥使用记录表
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS `key_usage` (
                `id` TEXT NOT NULL PRIMARY KEY,
                `key_id` TEXT NOT NULL,
                `operation` TEXT NOT NULL,
                `timestamp` TIMESTAMP NOT NULL,
                `user_id` TEXT NULL,
                `details` TEXT NULL,
                FOREIGN KEY (`key_id`) REFERENCES `keys` (`id`)
            );
            """)
            
            # 创建索引
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_key_id ON key_usage (key_id);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_timestamp ON key_usage (timestamp);")
            
            conn.commit()
        except:
            conn.rollback()
            raise
    
    def create_key(self, key_name, key_type="AES-256", key_data=None, expires_days=365, metadata=None):
        """
        创建并存储一个新密钥
        :param key_name: 密钥名称
        :param key_type: 密钥类型，默认AES-256
        :param key_data: 可选的自定义密钥数据，若为None则自动生成
        :param expires_days: 密钥过期天数，默认365天
        :param metadata: 可选的元数据
        :return: 密钥ID和生成时间(ms)
        """
        start_time = time.time()
        
        # 如果未提供密钥数据，则生成随机密钥
        if key_data is None:
            key_size = 32  # 256位
            key_data = os.urandom(key_size)
        
        # 生成唯一ID
        key_id = str(uuid.uuid4())
        now = datetime.now().isoformat()
        expires_at = (datetime.now() + timedelta(days=expires_days)).isoformat() if expires_days else None
        
        # 将密钥数据编码为base64
        key_data_b64 = base64.b64encode(key_data).decode('utf-8')
        
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("""
            INSERT INTO `keys` 
            (`id`, `key_name`, `key_type`, `key_data`, `salt`, `created_at`, `expires_at`, `metadata`) 
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                key_id, 
                key_name, 
                key_type, 
                key_data_b64,
                None,  # salt
                now,
                expires_at,
                json.dumps(metadata) if metadata else None
            ))
            conn.commit()
        except:
            conn.rollback()
            raise
        
        generation_time = (time.time() - start_time) * 1000  # 转换为毫秒
        return key_id, generation_time
    
    def create_derived_key(self, key_name, password, key_type="AES-256", expires_days=365, metadata=None):
        """
        基于密码创建并存储派生密钥
        :param key_name: 密钥名称
        :param password: 用于派生密钥的密码
        :param key_type: 密钥类型，默认AES-256
        :param expires_days: 密钥过期天数，默认365天
        :param metadata: 可选的元数据
        :return: 密钥ID和生成时间(ms)
        """
        start_time = time.time()
        
        # 生成随机盐值
        salt = os.urandom(16)
        
        # 使用PBKDF2派生密钥
        key_size = 32  # 256位
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=key_size,
            salt=salt,
            iterations=100000,
        )
        
        password_bytes = password.encode('utf-8') if isinstance(password, str) else password
        key_data = kdf.derive(password_bytes)
        
        # 生成唯一ID
        key_id = str(uuid.uuid4())
        now = datetime.now()
        expires_at = now + timedelta(days=expires_days) if expires_days else None
        
        # 将密钥数据和盐值编码为base64
        key_data_b64 = base64.b64encode(key_data).decode('utf-8')
        salt_b64 = base64.b64encode(salt).decode('utf-8')
        
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("""
            INSERT INTO `keys` 
            (`id`, `key_name`, `key_type`, `key_data`, `salt`, `created_at`, `expires_at`, `metadata`) 
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                key_id, 
                key_name, 
                key_type, 
                key_data_b64,
                salt_b64,
                now,
                expires_at,
                json.dumps(metadata) if metadata else None
            ))
            conn.commit()
        except:
            conn.rollback()
            raise
        
        generation_time = (time.time() - start_time) * 1000  # 转换为毫秒
        return key_id, generation_time
    
    def get_key(self, key_id=None, key_name=None):
        """
        获取密钥数据
        :param key_id: 密钥ID
        :param key_name: 密钥名称
        :return: 密钥数据（二进制）
        """
        if key_id is None and key_name is None:
            raise ValueError("必须提供key_id或key_name")
        
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            if key_id:
                cursor.execute("SELECT `key_data`, `revoked_at`, `expires_at` FROM `keys` WHERE `id` = ?", (key_id,))
            else:
                cursor.execute("SELECT `key_data`, `revoked_at`, `expires_at` FROM `keys` WHERE `key_name` = ? ORDER BY `created_at` DESC LIMIT 1", (key_name,))
            
            result = cursor.fetchone()
            if not result:
                return None
            
            # 将SQLite字符串日期转为datetime对象
            revoked_at = datetime.fromisoformat(result['revoked_at']) if result['revoked_at'] else None
            expires_at = datetime.fromisoformat(result['expires_at']) if result['expires_at'] else None
            
            # 检查密钥是否被撤销
            if revoked_at:
                raise ValueError(f"密钥已被撤销于 {revoked_at}")
            
            # 检查密钥是否过期
            if expires_at and expires_at < datetime.now():
                raise ValueError(f"密钥已过期于 {expires_at}")
                
            # 将Base64编码的密钥数据转换回二进制
            key_data = base64.b64decode(result['key_data'])
            
            # 记录密钥使用情况
            self._log_key_usage(key_id or key_name, "get_key")
            
            return key_data
        except Exception as e:
            raise e
    
    def revoke_key(self, key_id):
        """
        撤销密钥
        :param key_id: 密钥ID
        :return: 成功则返回True，失败返回False
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("UPDATE `keys` SET `revoked_at` = ? WHERE `id` = ?", 
                        (datetime.now().isoformat(), key_id))
            if cursor.rowcount == 0:
                return False
                
            conn.commit()
            self._log_key_usage(key_id, "revoke_key")
            return True
        except:
            conn.rollback()
            raise
    
    def rotate_key(self, old_key_id, new_key_data=None, expires_days=365):
        """
        轮换密钥
        :param old_key_id: 旧密钥ID
        :param new_key_data: 新密钥数据（可选），若为None则自动生成
        :param expires_days: 新密钥过期天数
        :return: 新密钥ID和轮换时间(ms)
        """
        start_time = time.time()
        
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            # 获取旧密钥信息
            cursor.execute("SELECT `key_name`, `key_type`, `metadata` FROM `keys` WHERE `id` = ?", (old_key_id,))
            result = cursor.fetchone()
            if not result:
                raise ValueError(f"找不到ID为 {old_key_id} 的密钥")
            
            # 创建新密钥，使用相同的名称和元数据，但附加轮换信息
            key_name = result['key_name']
            key_type = result['key_type']
            metadata = json.loads(result['metadata']) if result['metadata'] else {}
            
            # 添加轮换信息到元数据
            if 'rotation_history' not in metadata:
                metadata['rotation_history'] = []
            
            metadata['rotation_history'].append({
                'rotated_from': old_key_id,
                'rotated_at': datetime.now().isoformat()
            })
            
            # 创建新密钥
            new_key_id, _ = self.create_key(
                key_name=key_name,
                key_type=key_type,
                key_data=new_key_data,
                expires_days=expires_days,
                metadata=metadata
            )
            
            # 撤销旧密钥
            self.revoke_key(old_key_id)
            
            rotation_time = (time.time() - start_time) * 1000  # 转换为毫秒
            
            self._log_key_usage(new_key_id, "rotate_key", 
                             details=f"Rotated from key {old_key_id}")
            
            return new_key_id, rotation_time
        except:
            conn.rollback()
            raise
            
    def list_keys(self, include_revoked=False, include_expired=False):
        """
        列出所有密钥
        :param include_revoked: 是否包含已撤销的密钥
        :param include_expired: 是否包含已过期的密钥
        :return: 密钥列表，每个元素包含id、name、type、created_at等信息
        """
        conn = self._get_connection()
        now = datetime.now()
        
        try:
            cursor = conn.cursor()
            query = """
            SELECT `id`, `key_name`, `key_type`, `created_at`, `expires_at`, `revoked_at` 
            FROM `keys` WHERE 1=1
            """
            
            if not include_revoked:
                query += " AND `revoked_at` IS NULL"
            
            if not include_expired:
                query += " AND (`expires_at` IS NULL OR `expires_at` > ?)"
                cursor.execute(query, (now,))
            else:
                cursor.execute(query)
            
            keys = []
            for row in cursor.fetchall():
                # 将SQLite字符串日期转为datetime对象
                created_at = datetime.fromisoformat(row['created_at']) if row['created_at'] else None
                expires_at = datetime.fromisoformat(row['expires_at']) if row['expires_at'] else None
                revoked_at = datetime.fromisoformat(row['revoked_at']) if row['revoked_at'] else None
                
                key_info = {
                    'id': row['id'],
                    'name': row['key_name'],
                    'type': row['key_type'],
                    'created_at': created_at,
                    'expires_at': expires_at,
                    'revoked_at': revoked_at,
                    'is_active': revoked_at is None and 
                                (expires_at is None or expires_at > now)
                }
                keys.append(key_info)
            
            return keys
        except Exception as e:
            raise e
    
    def _log_key_usage(self, key_id, operation, details=None):
        """
        记录密钥使用情况
        :param key_id: 密钥ID
        :param operation: 操作类型
        :param details: 操作详情
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("""
            INSERT INTO `key_usage` (`id`, `key_id`, `operation`, `timestamp`, `details`)
            VALUES (?, ?, ?, ?, ?)
            """, (
                str(uuid.uuid4()),
                key_id,
                operation,
                datetime.now().isoformat(),
                details
            ))
            conn.commit()
        except:
            conn.rollback()
    
    def close(self):
        """关闭数据库连接"""
        if self._connection:
            self._connection.close()
            self._connection = None 